from matplotlib import pyplot
import math
import cmath
import DynamicWindowApproach as dyn
import dwa

class Point:
    def __init__(self):
        self.x = 0
        self.y = 0
        self.z = 0

###########################################################
#
#       用laser数据导航时必须将scan数据融合到同一坐标系下
#       定位车轮时可以考虑单独使用每个激光雷达的相对坐标系
#
###########################################################

def init_points(s, threshold, max_range=10.0):
    '''
    说明
    =======================================
    初始化雷达数据点，丢弃强度小于threshold, 和距离大于max_range的数据

        - param s
        雷达数据

        - param threshold
        强度阈值

        - param max_range
        距离阈值
    '''

    new_list = list()

    for i in range(len(s['ranges'])):
        s['angles'].append(s['angle_min']+i*math.pi/540) #角度间隔是0.33度，所以除以540
        if s['intensities'][i] >= threshold and s['ranges'][i] < max_range:
            new_list.append(i)

    new_ranges = list()
    new_angles = list()

    for i in new_list:
        new_ranges.append(s['ranges'][i])
        new_angles.append(s['angles'][i])

    print("data length change from %d to %d"%(len(s['ranges']), len(new_list)))

    s['angles'] = new_angles
    s['ranges'] = new_ranges


def transform(s, dr, dth):
    '''
    说明
    =======================================
    将雷达数据进行坐标变换，先平移，再旋转

        - param s
        雷达数据

        - param dr
        平移距离

        - param dth
        旋转角度，逆时针为正方向
    '''

    def reduce_ang(ang):
        '''
        将角度限制在-pi ~ pi
        '''
        while abs(ang)>math.pi:
            if ang>0: ang=ang-math.pi*2 
            else: ang=ang+math.pi*2

    a = s['angles']

    for i in range(len(a)):
        reduce_ang(a[i])

    sout = {
        "angles": [],
        "ranges": []
    }

    for i in range(len(s['ranges'])):
        sout['ranges'].append(math.sqrt((s['ranges'][i]*math.cos(a[i])+dr)**2+(s['ranges'][i]*math.sin(a[i]))**2))
        # 用正弦定理求新的角度
        if sout['ranges'][i] == 0:
            _a = 0
        else:
            _a = math.asin(s['ranges'][i]*math.sin(a[i])/sout['ranges'][i])

        # 解决arcsin求得角度范围限制在-pi/2 ~ pi/2的问题
        if a[i] > math.pi/2:
            if s['ranges'][i]*math.cos(a[i])+dr < 0:
                _a = math.pi - _a
        elif a[i] < -math.pi/2:
            if s['ranges'][i]*math.cos(a[i])+dr < 0:
                _a = -math.pi - _a
        sout['angles'].append(_a)

    sout['angles'] = list(map(lambda x:x+dth, sout['angles']))

    return sout

class ScanBase:
    @staticmethod
    def newInstance():
        return {
            "angles": [],
            "ranges": []
        }

def gather_scan(scans):
    '''
    说明
    =======================================
    将数据变换后的前后左右四个激光雷达数据整合成一个数据

        - param scan_t
        [front, left, back, right]雷达数据
    '''
    scan_t = ScanBase.newInstance()
    scan = ScanBase.newInstance()

    for s in scans:
        scan_t['angles'].extend(s['angles'])
        scan_t['ranges'].extend(s['ranges'])

    x = np.array(scan_t['angles'])
    si = x.argsort() #提取排序的索引

    for i in range(len(si)):
        scan['ranges'].append(scan_t['ranges'][i])
        scan['angles'].append(scan_t['angles'][i])

    return scan

def plot_scan(s, color, clear=True):
    if clear:
        pyplot.clf()
    # pyplot.polar(s['angles'], s['ranges'], color, markersize=3)

if __name__ == '__main__':
    import LaserData
    import numpy as np
    thre = 50
    buick_scan = None

    def init_four_lasers():
        '''四个激光雷达数据融合'''
        init_points(LaserData.front, thre)
        init_points(LaserData.left, thre)
        init_points(LaserData.right, thre)
        init_points(LaserData.back, thre)

        scan_f = transform(LaserData.front, 0.908, 0)
        scan_l = transform(LaserData.left, 0.518, math.pi/2)
        scan_r = transform(LaserData.right, 0.518, -math.pi/2)
        scan_b = transform(LaserData.back, 0.908, math.pi)

        pyplot.clf()

        pyplot.polar(scan_f['angles'], scan_f['ranges'], 'r.', markersize=1)
        pyplot.polar(scan_l['angles'], scan_l['ranges'], 'y.', markersize=1)
        pyplot.polar(scan_r['angles'], scan_r['ranges'], 'b.', markersize=1)
        pyplot.polar(scan_b['angles'], scan_b['ranges'], 'g.', markersize=1)
        ax = pyplot.gca(projection='polar')

        pyplot.pause(0.5)

        scans = [scan_f, scan_l, scan_b, scan_r]
        scan = gather_scan(scans)

    def init_buick_scan():
        init_points(LaserData.front_buick, thre, 3)
        scan = transform(LaserData.front_buick, 0.908, 0)

        # pyplot.polar(scan['angles'], scan['ranges'], 'r.', markersize=1)
        # ax = pyplot.gca(projection='polar')
        # pyplot.pause(0.5)

        return scan

    def detect_success(points, wheel_num=2):
        ret = True

        if wheel_num == 2:
            cal_couples = [(0,1)]
        else:
            cal_couples = [(0,1),(0,2),(0,3),(1,2),(1,3),(2,3)]
        
        def distance(p1, p2):
            return math.sqrt((p1.x-p2.x)**2 + (p1.y-p2.y)**2)

        for d in list(map(lambda c: distance(points[c[0]],points[c[1]]), cal_couples)):
            if d < 1.2:
                ret = False

        for d  in list(map(lambda p: math.sqrt(p.x**2+p.y**2), points)):
            if d < 0.1:
                ret = False

        return ret
        
    def detect(s, wheel_num=2):
        '''
        车轮检测
        =============

        - param s
            激光雷达数据

        - param wheel_num
            扫描车轮的数量，如果是2，则找到一对轮子就返回，如果是4，则必须找到两对轮子
        '''

        back = LaserData.wheel_edge_detect(s)

        r = [
            s['ranges'][back[0]],
            s['ranges'][back[1]]
        ]

        theta = [
                s['angles'][back[0]],
                s['angles'][back[1]]
        ]

        if wheel_num == 2:
            return {
                "r" : r,
                "theta" : theta,
                "ok": False
            }

        back_r = r[0]
        half_wheelbase = 1.8/2  #一半的轴距

        # 调整搜索阈值，前轮比后轮更近
        front_distance_threshold = back_r - half_wheelbase/math.cos(theta[0])
        front = LaserData.wheel_edge_detect(s, front_distance_threshold)

        r = [
                s['ranges'][back[0]],
                s['ranges'][back[1]],
                s['ranges'][front[0]],
                s['ranges'][front[1]]
            ]

        theta = [
                s['angles'][back[0]],
                s['angles'][back[1]],
                s['angles'][front[0]],
                s['angles'][front[1]]
        ]

        return {
            "r" : r,
            "theta" : theta,
            "ok": False
        }

    def pollar2cartesian(scan_pollar):
        '''
        极坐标转直角坐标
        ===============

        - param scan_pollar
          (r, phi)极坐标点
        '''
        cn = cmath.rect(*scan_pollar)
        return (cn.real, cn.imag)


    def get_marker_points(wheel_inf):
        points = list()
        
        for i in range(len(wheel_inf['r'])):
            cn = cmath.rect(wheel_inf['r'][i], wheel_inf['theta'][i])
            p = Point()
            p.x = cn.real
            p.y = cn.imag
            p.z = 0
            points.append(p)

        return points

    buick_scan = init_buick_scan()
    wheels = detect(buick_scan)

    if detect_success(get_marker_points(wheels)):
        # pyplot.polar(wheels['theta'], wheels['r'], 'c', linewidth=2)
        # pyplot.pause(0.5)

        p1 = LaserData.pollar2cartesian(wheels['r'][0], wheels['theta'][0])
        p2 = LaserData.pollar2cartesian(wheels['r'][1], wheels['theta'][1])

        gx = (p1[0]+p2[0])/2 + 1.0
        gy = (p1[1]+p2[1])/2

        dwa.main(gx, gy, LaserData.get_obstacle_array(buick_scan), dwa.RobotType.rectangle)

        # DWA = dyn.Core()
        # DWA.set_goal(gx, gy)
        # DWA.update_ob(buick_scan)
        # DWA.run_task(sim=True)



